import onnx
import yaml, sys
import numpy as np
from yaml.loader import SafeLoader
from . import utils

def parse_yaml(yaml_file):
    hw_list = []
    std_list = []
    mean_list = []
    resize_list = []
    crop_list = []
    control_list = []

    # Open the file and load the file
    with open(yaml_file) as f:
        data = yaml.load(f, Loader=SafeLoader)
        if 'norm' in data.keys():
            #print('got norm---')
            norm_list = data['norm']
            for n in norm_list:
                #print('n:', n)
                if 'std' in n.keys():
                    std_list_ = n['std']
                    if len(std_list_) == 3 or len(std_list_) == 1:
                        for n in std_list_:
                            if n > 0.0:
                                std_list.append(1.0/n)
                                print('add std:', 1.0/n)
                            else:
                                std_list.append(1.0/1e-6)    

                        #print('got std values:', std_list)

                    continue    

                if 'mean' in n.keys():
                    mean_list_ = n['mean']
                    if len(mean_list_) == 3 or len(mean_list_) == 1:
                        for n in mean_list_:
                            mean_list.append(int(n))

                        #print('got mean values:', mean_list) 

                    continue         

        if 'preproc' in data.keys():
            #print('got preproc---')
            preproc_list = data['preproc']
            for p in preproc_list:
                #print('p:',p)
                if 'hw' in p.keys():
                    hw_list_ = p['hw']
                    if len(hw_list_) == 2:
                        hw_list = hw_list_
                        #print('got hw values:', hw_list)
                    continue    

                if 'resize' in p.keys():
                    resize_list_ = p['resize']
                    if len(resize_list_) == 2 or len(resize_list_) == 1:
                        resize_list = resize_list_
                        #print('got resize values:', resize_list)
                    continue

                if 'crop' in p.keys():
                    crop_list_ = p['crop']
                    if len(crop_list_) == 4:
                        crop_list = crop_list_
                        #print('got crop values:', crop_list)
                    continue

                if 'control' in p.keys():
                    control_list_ = p['control']
                    if len(control_list_) == 4 or len(control_list_) == 5:
                        control_list = control_list_
                        #print('got control values:', control_list)
                    continue

    if len(std_list) == 0 or len(mean_list) == 0 or len(std_list) != len(mean_list):
        return {}

    preproc_dict = {} 

    preproc_dict['std'] = std_list
    preproc_dict['mean'] = mean_list

    if len(hw_list) == 0:
        hw_list = [-1, -1]

    preproc_dict['hw'] = hw_list

    if len(resize_list) == 0:
        resize_list=[-1, -1]

    preproc_dict['resize'] = resize_list

    if len(crop_list) == 0:
        crop_list = [-1, -1, -1, -1]

    preproc_dict['crop'] = crop_list

    if len(control_list) == 0:
        control_list = [False, True, False, False]

    preproc_dict['control'] = control_list

    return preproc_dict    

def get_const_type(val):
    types = {
        'const_std_r' : 1,
        'const_std_g' : 1,
        'const_std_b' : 1,
        'const_mean_r' : 2,
        'const_mean_g' : 2,
        'const_mean_b' : 2,
        'const_resize' : 6,
        'const_crop' : 6,
        'const_control' : 6
    }

    return types.get(val, None)

def get_init_val(init):
    val = []

    if init.data_type != get_const_type(init.name):
        print('ERROR, mismatch init type', get_const_type(init.name), init.data_type, init.name)
        return False

    dtype = init.data_type
    np_dtype = utils.convert_ort_type_2_np(dtype)
    if init.raw_data:
        params_list = np.fromstring(init.raw_data, dtype=np_dtype)
        for p in params_list:
            val.append(p)
    else:
        data_list = utils.get_data_list(dtype, init)
        for p in data_list:
            val.append(p)

    return val        

def get_snr(gpu_array, cpu_array):
    diff_array = np.subtract(cpu_array, gpu_array)
    x = np.square(diff_array)
    x = np.sum(x)

    y = np.square(cpu_array)
    y = np.sum(y) 

    snr = (x) / (y + 1e-7)

    snr = np.mean(snr)

    #print('snr:', snr)

    return snr  

def analyze(onnxfile, **kwargs):
    yaml = kwargs['yaml']
    origin_model_path = kwargs['origin_model']

    preproc_dict = parse_yaml(yaml) 

    model = onnx.load(onnxfile)

    input_ = model.graph.input[0]
    if input_.type.tensor_type.elem_type != 2:
        print('Input type error:', input_type)

    model_input_hw = [input_.type.tensor_type.shape.dim[-2].dim_value, input_.type.tensor_type.shape.dim[-1].dim_value]

    print('model_input_hw:', model_input_hw)

    expect_input = ['const_std_r', 'const_std_g', 'const_std_b','const_mean_r', 'const_mean_g','const_mean_b','const_resize', 'const_crop', 'const_control']

    first_check = False
    for id, node in enumerate(model.graph.node):
        if node.op_type == 'PreProc':
            #print('Got PreProc Node, input:', node.input)
            if set(expect_input).issubset(set(node.input)):
                #print('fist check is ok, id = ', id)
                first_check = True
                break
            else:
                return False

    if first_check == True and id == 0:
        init_list = []
        for init in model.graph.initializer:
            init_list.append(init.name)

        if set(expect_input).issubset(set(init_list)) == False:
            print('Error: expect_input is not in init_list')
            return False

        d = {}
        for init in model.graph.initializer:
            if init.name in expect_input:
                val = get_init_val(init)
                #print('get val:', init.name, val)
                d[init.name] = val

        for k, v in preproc_dict.items():
                key = 'const_' + k
                if key in d.keys() and k != 'resize':
                    #print('kv', k, ':', v)
                    np_array_1 = np.array(v).flatten()
                    np_array_2 = np.array(d[key]).flatten()
                    snr = get_snr(np_array_1, np_array_2)
                    if snr > 0.1:
                        print('ERROR:snr too large')
                        return False

        #print('d[const_resize]:', d['const_resize'])
        print('preproc_dict[hw]:', preproc_dict['hw'])

        if model_input_hw != preproc_dict['hw']:
            print('Error, model_input_hw != preproc_dict[\'hw\']', model_input_hw, preproc_dict['hw'])
            return False

        if len(d['const_resize']) != 4 or d['const_resize'][:2] != preproc_dict['resize'][:2]:
            print('ERROR: rezie param')
            return False

        model = onnx.load(origin_model_path)
        input_ = model.graph.input[0]
        input_hw = [input_.type.tensor_type.shape.dim[-2].dim_value, input_.type.tensor_type.shape.dim[-1].dim_value]

        if d['const_resize'][2:] != input_hw:
            print('Error: input_hw is wrong', input_hw, d['const_resize'][2:])
            return False

        return True    

    return False